package org.aeon.fasttasks.core.manager;

import org.aeon.fasttasks.core.Task;
import org.aeon.fasttasks.core.annotation.*;
import org.aeon.fasttasks.core.recorder.AsyncCsvRecorder;
import org.aeon.fasttasks.core.recorder.CsvRecorder;
import org.aeon.fasttasks.core.recorder.TaskRecorder;
import org.springframework.context.ApplicationContext;
import org.springframework.core.annotation.AnnotatedElementUtils;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;

import java.io.Closeable;
import java.lang.annotation.Annotation;
import java.lang.reflect.Field;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * @author aeon
 * @createTime 2025/08/05  13:54
 * @description
 */
public class TaskExpansionFactory {

    /**
     * 创建 TaskExpansion
     *
     * @param task
     * @param applicationContext
     * @param taskRecorder
     * @return
     */
    public static TaskExpansion create(Task task,
                                       ApplicationContext applicationContext,
                                       TaskRecorder taskRecorder) {
        Map<Class<? extends Annotation>, Annotation> annotations = getAnnotations(
                task.getClass(),
                EnableTask.class,
                DependentTasks.class,
                SkipTask.class,
                LoopTask.class,
                LoopTaskPro.class
        );

        proxyLoopTaskPro(task.getClass(), annotations);

        List<Closeable> closeables = new ArrayList<>();

        injectAsyncCsvRecorder(task, applicationContext, closeables);

        injectCsvRecorder(task);

        TaskContext taskContext = new TaskContext(
                task,
                Map.copyOf(annotations),
                List.copyOf(closeables),
                annotations.containsKey(LoopTask.class) ||
                        annotations.containsKey(LoopTaskPro.class)
        );

        return injectOther(new TaskExpansion(taskContext, taskRecorder), applicationContext);
    }

    /**
     * 克隆
     *
     * @param taskExpansion
     * @return
     */
    public static TaskExpansion clone(TaskExpansion taskExpansion,
                                      TaskRecorder taskRecorder) {
        TaskContext taskContext = taskExpansion.getTaskContext();
        TaskContext cloneContext = new TaskContext(
                taskContext.getTaskCounter(),
                taskContext.getTask(),
                taskContext.getAnnotations(),
                taskContext.getCloseables(),
                taskContext.getMultiStateLatch(),
                taskContext.getTaskData());
        cloneContext.setBatchWorkTasksExecutor(taskContext.getBatchWorkTasksExecutor());
        cloneContext.setBossTasksExecutor(taskContext.getBossTasksExecutor());
        cloneContext.setTasksScheduler(taskContext.getTasksScheduler());

        cloneContext.setFrontEndTasks(taskContext.getFrontEndTasks().toArray(new TaskExpansion[]{}));
        TaskExpansion cloneTask = new TaskExpansion(cloneContext, taskRecorder);

        cloneTask.setFuture(taskExpansion.getFuture());
        return cloneTask;
    }

    /**
     * 确保 LoopTaskPro 的 sync 和 async 注解的 async 值正确
     *
     * @param annotations
     */
    private static void proxyLoopTaskPro(Class<?> clazz, Map<Class<? extends Annotation>, Annotation> annotations) {
        if (annotations.containsKey(LoopTaskPro.class)) {
            LoopTaskPro loopTaskPro = (LoopTaskPro) annotations.get(LoopTaskPro.class);
            if (loopTaskPro.value().length == 0) {
                Map<String, Object> syncAttributes = AnnotationUtils.getAnnotationAttributes(loopTaskPro.sync());
                syncAttributes.put("async", false);
                Map<String, Object> asyncAttributes = AnnotationUtils.getAnnotationAttributes(loopTaskPro.async());
                asyncAttributes.put("async", true);

                annotations.put(LoopTaskPro.class, AnnotationUtils.synthesizeAnnotation(Map.of(
                        "sync", AnnotationUtils.synthesizeAnnotation(syncAttributes, LoopTask.class, null),
                        "async", AnnotationUtils.synthesizeAnnotation(asyncAttributes, LoopTask.class, null)
                ), LoopTaskPro.class, clazz));
            } else if (loopTaskPro.value().length == 2 &&
                    (loopTaskPro.value()[0].async() ^ loopTaskPro.value()[1].async())) {
                LoopTask syncTask = loopTaskPro.value()[0].async() ? loopTaskPro.value()[1] : loopTaskPro.value()[0];
                LoopTask asyncTask = loopTaskPro.value()[0].async() ? loopTaskPro.value()[0] : loopTaskPro.value()[1];

                annotations.put(LoopTaskPro.class, AnnotationUtils.synthesizeAnnotation(Map.of(
                        "sync", AnnotationUtils.synthesizeAnnotation(syncTask, LoopTask.class),
                        "async", AnnotationUtils.synthesizeAnnotation(asyncTask, LoopTask.class)
                ), LoopTaskPro.class, clazz));
            } else {
                throw new IllegalArgumentException("LoopTaskPro.value() 参数错误");
            }
        } else if (annotations.containsKey(LoopTask.class)) {
            Set<LoopTask> mergedRepeatableAnnotations =
                    AnnotatedElementUtils.getMergedRepeatableAnnotations(clazz, LoopTask.class, LoopTaskPro.class);
            if (mergedRepeatableAnnotations.size() == 1) {
                return;
            }
            LoopTask syncTask = null;
            LoopTask asyncTask = null;
            // 拿第一个满足要求的
            for (LoopTask next : mergedRepeatableAnnotations) {
                if (next.async()) {
                    if (Objects.isNull(asyncTask)) {
                        asyncTask = next;
                    }
                } else {
                    if (Objects.isNull(syncTask)) {
                        syncTask = next;
                    }
                }
                if (Objects.nonNull(syncTask) && Objects.nonNull(asyncTask)) {
                    break;
                }
            }
            if (Objects.isNull(syncTask) || Objects.isNull(asyncTask)) {
                throw new IllegalArgumentException(clazz.getSimpleName() + " 类上的 @LoopTask 注释错误");
            }
            annotations.put(LoopTaskPro.class, AnnotationUtils.synthesizeAnnotation(Map.of(
                    "sync", syncTask,
                    "async", asyncTask
            ), LoopTaskPro.class, clazz));
        }
    }

    /**
     * 注入 AsyncCsvRecorder
     *
     * @param task
     * @param applicationContext
     * @param closeables
     */
    private static void injectAsyncCsvRecorder(Task task,
                                               ApplicationContext applicationContext,
                                               List<Closeable> closeables) {
        closeables.addAll(getFieldsWithAnnotation(task.getClass(), InjectAsyncCsvRecorder.class)
                .entrySet()
                .stream()
                .map(entry -> {
                    // 注入 AsyncCsvLogger
                    InjectAsyncCsvRecorder annotation = (InjectAsyncCsvRecorder) entry.getKey();

                    AsyncCsvRecorder asyncCsvRecorder = new AsyncCsvRecorder(annotation.path(),
                            List.of(annotation.headers()), annotation.queueSize(),
                            annotation.batchSize(), annotation.recordingTime(),
                            applicationContext.getBean("bossTasksExecutor", ThreadPoolTaskExecutor.class));
                    try {
                        if (entry.getValue().getType().isAssignableFrom(CsvRecorder.class)) {
                            entry.getValue().setAccessible(true);
                            entry.getValue().set(task, asyncCsvRecorder);
                        } else {
                            throw new IllegalArgumentException(task.getClass().getName() +
                                    "." + entry.getValue().getName() + " 必须是 CsvRecorders 或其父类");
                        }
                    } catch (IllegalAccessException e) {
                        throw new RuntimeException(e);
                    }
                    return asyncCsvRecorder;
                })
                .toList());
    }

    /**
     * 注入 CsvRecorder
     *
     * @param task
     */
    private static void injectCsvRecorder(Task task) {
        getFieldsWithAnnotation(task.getClass(), InjectCsvRecorder.class)
                .forEach((k, v) -> {
                    // 注入 CsvLogger
                    InjectCsvRecorder annotation = (InjectCsvRecorder) k;
                    CsvRecorder csvRecorder = new CsvRecorder(annotation.path(),
                            List.of(annotation.headers()), annotation.recordingTime());
                    try {
                        if (v.getType().isAssignableFrom(CsvRecorder.class)) {
                            v.setAccessible(true);
                            v.set(task, csvRecorder);
                        } else {
                            throw new IllegalArgumentException(
                                    task.getClass().getName() + "." + v.getName() + " 必须是 CsvRecorders 或其父类");
                        }
                    } catch (IllegalAccessException e) {
                        throw new RuntimeException(e);
                    }
                });
    }

    /**
     * 其他注入
     *
     * @param task
     * @param applicationContext
     */
    private static TaskExpansion injectOther(TaskExpansion task, ApplicationContext applicationContext) {
        TaskContext taskContext = task.getTaskContext();
        taskContext.setBatchWorkTasksExecutor(applicationContext.getBean("batchWorkTasksExecutor",
                ThreadPoolTaskExecutor.class));
        taskContext.setBossTasksExecutor(applicationContext.getBean("bossTasksExecutor",
                ThreadPoolTaskExecutor.class));
        taskContext.setTasksScheduler(applicationContext.getBean("tasksScheduler",
                ThreadPoolTaskScheduler.class));
        return task;
    }


    /**
     * 从目标类上提取指定的注解，并返回不可变的 Map（键：注解类型，值：注解实例）
     *
     * @param clazz           目标类
     * @param annotationClass 要提取的注解类型
     * @return 不可变的 Map，仅包含存在的注解
     */
    @SafeVarargs
    private static Map<Class<? extends Annotation>, Annotation> getAnnotations(
            Class<?> clazz,
            Class<? extends Annotation>... annotationClass
    ) {
        return Stream.of(annotationClass)
                .map(annotationType -> new AbstractMap.SimpleEntry<>(
                        // 通过spring工具获取注解会处理AliasFor等
                        annotationType, AnnotationUtils.findAnnotation(clazz, annotationType)))
                .filter(entry -> entry.getValue() != null)
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
    }

    /**
     * 从目标类上提取指定注解的字段
     *
     * @param clazz
     * @param annotationClass
     * @return
     */
    private static Map<? extends Annotation, Field> getFieldsWithAnnotation(
            Class<?> clazz,
            Class<? extends Annotation> annotationClass
    ) {
        return Arrays.stream(clazz.getDeclaredFields())
                .filter(field -> field.isAnnotationPresent(annotationClass))
                .collect(Collectors.toMap(
                        f -> AnnotationUtils.findAnnotation(f, annotationClass),
                        field -> field
                ));
    }

}
